import pymysql
import logging
import time
from typing import List, Tuple, Optional
from datetime import datetime, date
from db_config import DBConfig  # 导入统一的DBConfig

# 导入分红同步函数
from stock_dividend import (
    sync_stock_dividend_to_db,
    batch_sync_stock_dividend
)

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class StockDividendTrigger:
    def __init__(self):
        self.db_config = DBConfig.config  # 使用统一的DBConfig
        self.conn = None
        self.connect_db()

    def connect_db(self):
        """连接数据库"""
        try:
            self.conn = pymysql.connect(
                host=self.db_config['host'],
                port=self.db_config['port'],
                user=self.db_config['user'],
                password=self.db_config['password'],
                database=self.db_config['database'],
                charset=self.db_config['charset']
            )
            logger.info("数据库连接成功")
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            raise

    def get_stocks_for_full_sync(self, query_sql: str) -> List[Tuple[str, str]]:
        """
        获取需要全量同步的股票列表

        :param query_sql: 查询SQL语句，需要返回stock_code和stock_name字段
        :return: 股票代码和名称列表 [(code, name), ...]
        """
        stocks = []
        try:
            with self.conn.cursor() as cursor:
                cursor.execute(query_sql)
                results = cursor.fetchall()

                for row in results:
                    if len(row) >= 2:
                        stock_code = row[0]  # stock_code字段
                        stock_name = row[1]  # stock_name字段

                        # stock_dividend.py使用纯数字代码，无需添加交易所后缀
                        stocks.append((stock_code, stock_name))
                    else:
                        logger.warning(f"查询结果格式不正确: {row}")
                        continue

                logger.info(f"全量同步查询到 {len(stocks)} 只股票")
                return stocks

        except Exception as e:
            logger.error(f"执行全量同步查询SQL失败: {e}")
            return []

    def get_stocks_for_incremental_sync(self, query_sql: str) -> List[Tuple[str, str]]:
        """
        获取需要增量同步的股票列表（当年未分红的股票）

        :param query_sql: 基础查询SQL语句，需要返回stock_code和stock_name字段
        :return: 股票代码和名称列表 [(code, name), ...]
        """
        current_year = datetime.now().year
        stocks = []

        try:
            with self.conn.cursor() as cursor:
                # 查询当年未分红的股票
                incremental_sql = f"""
                SELECT sb.stock_code, sb.stock_name
                FROM ({query_sql}) sb
                LEFT JOIN finance_dividend fd ON sb.stock_code = fd.stock_code
                    AND YEAR(fd.ex_dividend_date) = {current_year}
                WHERE fd.stock_code IS NULL
                """

                cursor.execute(incremental_sql)
                results = cursor.fetchall()

                for row in results:
                    if len(row) >= 2:
                        stock_code = row[0]  # stock_code字段
                        stock_name = row[1]  # stock_name字段

                        # stock_dividend.py使用纯数字代码，无需添加交易所后缀
                        stocks.append((stock_code, stock_name))
                    else:
                        logger.warning(f"查询结果格式不正确: {row}")
                        continue

                logger.info(f"增量同步查询到 {len(stocks)} 只股票（当年未分红）")
                return stocks

        except Exception as e:
            logger.error(f"执行增量同步查询SQL失败: {e}")
            return []

    def trigger_dividend_sync(self,
                            query_sql: str,
                            sync_type: str = "incremental",
                            batch_size: int = 10,
                            delay_seconds: int = 2):
        """
        触发股票分红数据同步

        :param query_sql: 查询SQL语句，用于获取需要同步的股票代码和名称
        :param sync_type: 同步类型，"full"表示全量同步，"incremental"表示增量同步
        :param batch_size: 每批处理的股票数量
        :param delay_seconds: 每批股票同步后的延迟时间（秒）
        """
        if sync_type == "full":
            logger.info("开始全量同步分红数据")
            stock_list = self.get_stocks_for_full_sync(query_sql)
        elif sync_type == "incremental":
            logger.info("开始增量同步分红数据（当年未分红的股票）")
            stock_list = self.get_stocks_for_incremental_sync(query_sql)
        else:
            logger.error(f"不支持的同步类型: {sync_type}")
            return

        if not stock_list:
            logger.warning("没有找到需要同步的股票")
            return

        logger.info(f"开始同步 {len(stock_list)} 只股票的分红数据")

        # 分批处理
        for i in range(0, len(stock_list), batch_size):
            batch = stock_list[i:i + batch_size]
            logger.info(f"正在处理第 {i//batch_size + 1}/{(len(stock_list)-1)//batch_size + 1} 批，本批 {len(batch)} 只股票")

            # 调用分红同步函数
            batch_sync_stock_dividend(batch)

            # 如果不是最后一批，添加延迟
            if i + batch_size < len(stock_list):
                logger.info(f"等待 {delay_seconds} 秒后继续下一批...")
                time.sleep(delay_seconds)

        logger.info("所有股票分红数据同步完成")

    def get_dividend_sync_status(self) -> dict:
        """
        获取分红同步状态统计
        """
        try:
            with self.conn.cursor() as cursor:
                # 统计总股票数
                cursor.execute("SELECT COUNT(*) FROM stock_basic WHERE is_stock_st = 0")
                total_stocks = cursor.fetchone()[0]

                # 统计已同步分红数据的股票数
                current_year = datetime.now().year
                cursor.execute(f"""
                    SELECT COUNT(DISTINCT stock_code)
                    FROM finance_dividend
                    WHERE YEAR(ex_dividend_date) = {current_year}
                """)
                synced_stocks = cursor.fetchone()[0]

                # 统计分红记录总数
                cursor.execute("SELECT COUNT(*) FROM finance_dividend")
                total_records = cursor.fetchone()[0]

                return {
                    "total_stocks": total_stocks,
                    "synced_stocks": synced_stocks,
                    "unsynced_stocks": total_stocks - synced_stocks,
                    "total_records": total_records,
                    "sync_percentage": round((synced_stocks / total_stocks * 100), 2) if total_stocks > 0 else 0
                }

        except Exception as e:
            logger.error(f"获取同步状态失败: {e}")
            return {}

    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
            logger.info("数据库连接已关闭")

# 使用示例
if __name__ == "__main__":
    trigger = StockDividendTrigger()

    try:
        # 示例1：全量同步所有非ST股票
        full_sync_sql = """
        SELECT stock_code, stock_name
        FROM stock_basic
        WHERE is_stock_st = 0
        """

        # 示例2：增量同步（默认）
        incremental_sync_sql = """
        SELECT stock_code, stock_name
        FROM stock_basic
        WHERE is_stock_st = 0
        """

        # 示例3：同步特定行业的股票
        # industry_sync_sql = """
        # SELECT stock_code, stock_name
        # FROM stock_basic
        # WHERE industry_type LIKE '%银行%'
        # AND is_stock_st = 0
        # """

        # 示例4：同步特定市场的股票
        # market_sync_sql = """
        # SELECT stock_code, stock_name
        # FROM stock_basic
        # WHERE market_type = 'sh_main'
        # AND is_stock_st = 0
        # """

        # 注意：stock_dividend.py使用纯数字股票代码格式，如"002236"，无需交易所后缀

        # 执行全量同步
        print("=== 执行全量同步 ===")
        trigger.trigger_dividend_sync(
            query_sql=full_sync_sql,
            sync_type="full",
            batch_size=10,
            delay_seconds=2
        )

        # 执行增量同步
        # print("\n=== 执行增量同步 ===")
        # trigger.trigger_dividend_sync(
        #     query_sql=incremental_sync_sql,
        #     sync_type="incremental",
        #     batch_size=10,
        #     delay_seconds=2
        # )

        # 获取同步状态
        print("\n=== 同步状态统计 ===")
        status = trigger.get_dividend_sync_status()
        for key, value in status.items():
            print(f"{key}: {value}")

    except Exception as e:
        logger.error(f"触发同步过程中发生错误: {e}")

    finally:
        trigger.close()